import sys

import torch
from causal.model1 import Recommender  # 确保模型类已导入

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = Recommender(
    dim=64,
    num_user=1101,
    num_item=711,
    device=device
).to(device)


model.load_state_dict(torch.load('model.pkl'))
model.eval()


import pymysql

conn = pymysql.connect(host='39.101.70.41', port=3306, user='myapp', passwd='cDnS7KZZCjYnSZZ8', db='myapp',
                       charset='utf8')  # 打开数据库连接
cursor = conn.cursor()
sql = "select * from user_rating1"
cursor.execute(sql)
results = cursor.fetchall()

from collections import *

users = defaultdict(lambda: defaultdict(float))
for a, b, c in results:
    users[int(a)][b] = c
user_ids = list(users.keys())
id_users = []
user_ids = [int(i) for i in user_ids]
items_ids = []
session_items = []
recent_items = []
for user in user_ids:
    scores = list(users[user].items())
    scores.sort(key=lambda x: x[1], reverse=True)
    ids = [score[0] for score in scores]
    session_items.append(ids[:5])
    recent_items.append(ids[-5:])

user_ids = torch.LongTensor(user_ids)
session_items = torch.LongTensor(session_items)
recent_items = torch.LongTensor(recent_items)

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
user_ids = user_ids.to(device)
session_items = session_items.to(device)
recent_items = recent_items.to(device)

with torch.no_grad():
    # 获取预测分数
    scores, _, _, _, _, _, _ = model(
        user_ids,
        session_items[:, :-1],  # 输入会话上下文（排除最后一个物品）
        recent_items
    )

    # 对最后一个位置进行预测
    last_preds = scores[:, -1, :]  # [batch_size, num_items]
    topk_items = torch.topk(last_preds, k=4, dim=1)[1]

rec_items = topk_items.cpu().numpy()
results = []
for items in rec_items:
    item_t = [(799 - int(i)) // 20 for i in items]
    results.append(item_t)
for i in range(len(results)):
    for j in range(len(results[i])):
        while results[i].count(results[i][j]) > 1:
            results[i][j] = results[i][j] + 1

def main(id):
    return results[int(id)]

if __name__ == '__main__':
    id = sys.argv[1]
    print(main(id))